// "Copyright [2021] <Copyright Shengkai Lin>"

#include "ocs.h"

namespace ai_sim {

Define_Module(OCS);

using namespace omnetpp;

void OCS::initialize() {
    switch_delay = par("switch_delay");
    my_index = par("my_index");

    for (int i = 0; i < gateSize("optical_port$o"); i++) {
        opti_gate* self_gate = gate("optical_port$o", i);
        int self_gate_id = self_gate->getId();
        cGate* remote_gate = self_gate->getNextGate();
        opti_out_gates[self_gate_id] = remote_gate;
        EV << self_gate->getFullPath() << " push to opti_out_gates"<< endl;
        EV << "key is " << self_gate->getId() << endl;
    }
    for (int i = 0; i < gateSize("optical_port$i"); i++) {
        opti_gate* self_gate = gate("optical_port$i", i);
        int self_gate_id = self_gate->getId();
        cGate* remote_gate = self_gate->getPreviousGate();
        opti_in_gates[self_gate_id] = remote_gate;
        EV << self_gate->getFullPath() << " push to opti_in_gates"<< endl;
        EV << "key is " << self_gate->getId() << endl;
    }
    EV << "OCS init finished, the opti gates size is " <<
        gateSize("optical_port$o") << endl;
}

void OCS::handleMessage(omnetpp::cMessage *msg) {
    auto* ra_msg = dynamic_cast<RouteVecMsg*>(msg);
    auto* oxc_msg = dynamic_cast<OXCConfMsg*>(msg);
    if (msg->isSelfMessage()) {
        // The msg is a oxc_ack_msg, then send to controller
        // OXCConfMsg *oxc_msg = dynamic_cast<OXCConfMsg *> (msg);
        // for (auto connection : oxc_msg->getChange_set().config_to_add) {
        //     cChannelType *channelType =
        //         cChannelType::get("cluster.network_devices.Fiber");
        //     cChannel *new_channel = channelType->create("channel");
        //     opti_gate* inPort = opti_in_gates[connection.in_port_id];
        //     opti_gate* outPort = opti_out_gates[connection.out_port_id];
        //     if (inPort->isConnected()) {
        //         if (inPort->getNextGate() != outPort) {
        //             EV << "The gate is connected to " << inPort->getNextGate() << endl;
        //             EV << "Try to connect to " << outPort << endl;
        //             error("connectTo Error");
        //         }
        //     } else {
        //         inPort->connectTo(outPort, new_channel);
        //     }
        // }
        // SwitchAckMsg *oxc_ack_msg = new SwitchAckMsg;
        // oxc_ack_msg->setThe_flow(oxc_msg->getThe_flow());
        // oxc_ack_msg->setKind(SWITCH_ACK_MSG);
        // send(oxc_ack_msg, "controller_port$o");
        msg->setKind(SWITCH_ACK_MSG);
        send(msg, "controller_port$o");
        delete oxc_msg;
        delete ra_msg;
    } else if (msg->getKind() == ROUTE_ALLOC_MSG || ra_msg != nullptr) {
        EV << "start alloc routes by ROUTE_ALLOC_MSG" << endl;
        RouteVecMsg *ra_msg = dynamic_cast<RouteVecMsg *> (msg);
        assert(ra_msg->getTargetId() == getId());
        for (auto routes : ra_msg->getNew_route_table()) {
            // Switch immediately if config on start
            opti_gate* srcPort = nullptr;
            opti_gate* destPort = nullptr;

            if (opti_in_gates.find(routes.inPort) != opti_in_gates.end()) {
                srcPort = opti_in_gates[routes.inPort];
                EV << "Disconnecting " << srcPort->getFullPath() << endl;
                srcPort->disconnect();
            }
            if (opti_out_gates.find(routes.outPort) != opti_out_gates.end()) {
                destPort = opti_out_gates[routes.outPort];
                EV << "Disconnecting " << destPort->getFullPath() << endl;
                // disconnect the current connection
                if (destPort->isConnectedOutside()) {
                    destPort->getPreviousGate()->disconnect();
                }
            }

            if (srcPort == nullptr || destPort == nullptr) {
                cout << "require inPort id " << routes.inPort << endl;
                cout << "require outPort id " << routes.outPort << endl;
                error("inPort or outPort is nullptr");
            }
            cChannelType *channelType =
                cChannelType::get("cluster.network_devices.Fiber");
            cChannel *new_channel = channelType->create("channel");
            if (destPort->isConnectedOutside()) {
                cout << "Port " << destPort->getFullPath() <<" is connected " <<
                    destPort->getPreviousGate()->getFullPath() << endl;
                cout << "is trying to connect " <<
                    srcPort->getFullPath() << " to it" << endl;
                error("connectTo error");
            } else {
                EV << "Arrived" << endl;
                srcPort->connectTo(destPort, new_channel);
            }
        }
        delete ra_msg;

    } else if (msg->getKind() == OXC_CONF_MSG || oxc_msg != nullptr) {
        OXCConfMsg *oxc_msg = dynamic_cast<OXCConfMsg *> (msg);
        for (auto connection : oxc_msg->getChange_set().config_to_remove) {
            opti_gate* inPort = opti_in_gates[connection.in_port_id];
            opti_gate* outPort = opti_out_gates[connection.out_port_id];
            if (outPort->isConnectedOutside()) {
                outPort->getPreviousGate()->disconnect();
            }
            inPort->disconnect();
        }
        for (auto connection : oxc_msg->getChange_set().config_to_add) {
            cChannelType *channelType =
                cChannelType::get("cluster.network_devices.Fiber");
            cChannel *new_channel = channelType->create("channel");
            opti_gate* inPort = opti_in_gates[connection.in_port_id];
            opti_gate* outPort = opti_out_gates[connection.out_port_id];
            if (inPort->isConnected()) {
                if (inPort->getNextGate() != outPort) {
                    EV << "The gate is connected to " << inPort->getNextGate() << endl;
                    EV << "Try to connect to " << outPort << endl;
                    error("connectTo Error");
                }
            } else {
                inPort->connectTo(outPort, new_channel);
            }
        }
        SwitchAckMsg *oxc_ack_msg = new SwitchAckMsg;
        oxc_ack_msg->setThe_flow(oxc_msg->getThe_flow());
        oxc_ack_msg->setKind(SWITCH_ACK_MSG);
        if((oxc_msg->getChange_set().config_to_remove.size() != 0) 
            || (oxc_msg->getChange_set().config_to_add.size() != 0)) {
                scheduleAt(omnetpp::simTime() + switch_delay, oxc_ack_msg);
            }
        else {
            scheduleAt(omnetpp::simTime(), oxc_ack_msg);
        }
        delete oxc_msg;
    } else {
        EV << "msg kind " << msg->getKind() << endl;
        std::cout << "Msg arrived on " <<
            msg->getArrivalGate()->getFullPath() << endl;
        error("Unknown msg kind");
    }
}

}  // namespace ai_sim
